import numpy as np

from train_components.uni_train import UniPLMTrainer
import torch
from models.plm_models import T5MatchModel
from evalulators.evaluation import plm_evaluate_test_data, plm_evaluate_valid_data
import os
from sklearn.metrics import roc_curve, roc_auc_score, precision_recall_curve,auc

# mhc2
train_corpus_path = './data/train.csv'
valid_corpus_path = './data/valid.csv'
test_corpus_path = './data/test.csv'

pretrained_path = './pretrained'
MODEL = T5MatchModel
model_name = 'T5match'
# 学习率
lr = 0.01
num_epochs = 20
batch_size = 64
max_length = 49
criterion = None
save_path = './Results-model'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

uni_trainer = UniPLMTrainer(pretrained_path=pretrained_path, train_corpus_path=train_corpus_path,
                            valid_corpus_path=valid_corpus_path,
                            MODEL=MODEL, lr=lr, num_epochs=num_epochs,
                            batch_size=batch_size, max_length=max_length, criterion=criterion,
                            save_path=save_path, model_name=model_name, device=device)
# uni_trainer.training()

model = MODEL.from_pretrained(os.path.join(save_path, model_name))
model.to(device=device)


real, pred, prob, metrics = plm_evaluate_test_data(model=model, save_path=save_path, test_corpus_path=test_corpus_path,
                    device=device, model_name=model_name, average='binary')

precision,recall,_ = precision_recall_curve(real,prob)
print(auc(recall,precision))


